package ai.koog.prompt.executor.clients.openrouter.models

import ai.koog.prompt.executor.clients.openai.base.models.OpenAIBaseLLMRequest
import ai.koog.prompt.executor.clients.openai.base.models.OpenAIBaseLLMResponse
import ai.koog.prompt.executor.clients.openai.base.models.OpenAIBaseLLMStreamResponse
import ai.koog.prompt.executor.clients.openai.base.models.OpenAIMessage
import ai.koog.prompt.executor.clients.openai.base.models.OpenAIResponseFormat
import ai.koog.prompt.executor.clients.openai.base.models.OpenAIStaticContent
import ai.koog.prompt.executor.clients.openai.base.models.OpenAITool
import ai.koog.prompt.executor.clients.openai.base.models.OpenAIToolCall
import ai.koog.prompt.executor.clients.openai.base.models.OpenAIToolChoice
import ai.koog.prompt.executor.clients.openai.base.models.OpenAIUsage
import ai.koog.prompt.executor.clients.serialization.AdditionalPropertiesFlatteningSerializer
import kotlinx.serialization.SerialName
import kotlinx.serialization.Serializable
import kotlinx.serialization.json.JsonElement

/**
 * OpenRouter Chat Completions API Request
 * https://openrouter.ai/docs/api-reference
 *
 */
@Serializable
internal class OpenRouterChatCompletionRequest(
    val messages: List<OpenAIMessage> = emptyList(),
    val prompt: String? = null,
    override val model: String? = null,
    override val stream: Boolean? = null,
    override val temperature: Double? = null,
    val tools: List<OpenAITool>? = null,
    val toolChoice: OpenAIToolChoice? = null,
    override val topP: Double? = null,
    override val topLogprobs: Int? = null,
    val maxTokens: Int? = null,
    val frequencyPenalty: Double? = null,
    val presencePenalty: Double? = null,
    val responseFormat: OpenAIResponseFormat? = null,
    val stop: List<String>? = null,
    val logprobs: Boolean? = null,
    val seed: Int? = null,
    val topK: Int? = null,
    val repetitionPenalty: Double? = null,
    val logitBias: Map<Int, Double>? = null,
    val minP: Double? = null,
    val topA: Double? = null,
    val prediction: OpenAIStaticContent? = null,
    val transforms: List<String>? = null,
    val models: List<String>? = null,
    val route: String? = null,
    val provider: ProviderPreferences? = null,
    val user: String? = null,
    val additionalProperties: Map<String, JsonElement>? = null,
) : OpenAIBaseLLMRequest

/**
 * @property order List of provider slugs to try in order (e.g. ["anthropic", "openai"])
 * @property allowFallbacks Whether to allow backup providers when the primary is unavailable
 * @property requireParameters Only use providers that support all parameters in your request
 * @property dataCollection `allow` Control whether to use providers that may store data
 * @property only List of provider slugs to allow for this request
 * @property ignore List of provider slugs to skip for this request
 * @property quantizations List of quantization levels to filter by (e.g. ["int4", "int8"])
 * @property sort Sort providers by price or throughput. (e.g. "price" or "throughput")
 * @property maxPrice The maximum pricing you want to pay for this request
 *
 */
@Serializable
public class ProviderPreferences(
    public val order: List<String>? = null,
    public val allowFallbacks: Boolean? = null,
    public val requireParameters: Boolean? = null,
    public val dataCollection: String? = null,
    public val only: List<String>? = null,
    public val ignore: List<String>? = null,
    public val quantizations: List<String>? = null,
    public val sort: String? = null,
    public val maxPrice: Map<String, String>? = null
)

/**
 * Chat completion choice
 *
 * @property finishReason The reason the model stopped generating tokens.
 * This will be `stop` if the model hit a natural stop point or a provided stop sequence,
 * `length` if the maximum number of tokens specified in the request was reached,
 * `content_filter` if content was omitted due to a flag from our content filters,
 * `tool_calls` if the model called a tool, or `function_call` (deprecated) if the model called a function.
 * `error` if the model finishes the request with an error.
 * @property nativeFinishReason The raw finish_reason string returned by the model
 * @property message A chat completion message generated by the model.
 * @property error An error response structure typically used for conveying error details to the clients.
 *
 * See (CompletionsResponse Format)[https://openrouter.ai/docs/api-reference/overview#completionsresponse-format]
 */
@Serializable
public class OpenRouterChoice(
    public val finishReason: String? = null,
    public val nativeFinishReason: String? = null,
    public val message: OpenAIMessage,
    public val error: ErrorResponse? = null
)

/**
 * Chat completion choice
 *
 * @property finishReason The reason the model stopped generating tokens.
 * This will be `stop` if the model hit a natural stop point or a provided stop sequence,
 * `length` if the maximum number of tokens specified in the request was reached,
 * `content_filter` if content was omitted due to a flag from our content filters,
 * `tool_calls` if the model called a tool, or `function_call` (deprecated) if the model called a function.
 * `error` if the model finishes the request with an error.
 * @property nativeFinishReason The raw finish_reason string returned by the model
 * @property delta A chat completion delta generated by streamed model responses.
 * @property error An error response structure typically used for conveying error details to the clients.
 *
 * See (CompletionsResponse Format)[https://openrouter.ai/docs/api-reference/overview#completionsresponse-format]
 */
@Serializable
public class OpenRouterStreamChoice(
    public val finishReason: String? = null,
    public val nativeFinishReason: String? = null,
    public val delta: OpenRouterStreamDelta,
    public val error: ErrorResponse? = null
)

/**
 * @property content The contents of the chunk message.
 * @property role The role of the author of this message.
 * @property toolCalls The tool calls requested by the model.
 */
@Serializable
public class OpenRouterStreamDelta(
    public val content: String? = null,
    public val role: String? = null,
    public val toolCalls: List<OpenAIToolCall>? = null
)

/**
 * Represents an error response structure typically used for conveying error details to the clients.
 *
 * @property code The numeric code representing the error.
 * @property message A descriptive message providing details about the error.
 * @property metadata Optional additional information about the error in the form of key-value pairs.
 */
@Serializable
public class ErrorResponse(
    public val code: Int,
    public val message: String,
    public val metadata: Map<String, String>? = null,
)

/**
 * OpenRouter Chat Completion Response
 * See (CompletionsResponse Format)[https://openrouter.ai/docs/api-reference/overview#completionsresponse-format]
 */
@Serializable
public class OpenRouterChatCompletionResponse(
    public val choices: List<OpenRouterChoice> = emptyList(),
    override val created: Long = 0L,
    override val id: String = "",
    override val model: String = "",
    public val systemFingerprint: String? = null,
    @SerialName("object")
    public val objectType: String = "chat.completion",
    public val usage: OpenAIUsage? = null,
    public val error: OpenRouterError? = null,
) : OpenAIBaseLLMResponse

/**
 * OpenRouter Chat Completion Streaming Response
 * See (CompletionsResponse Format)[https://openrouter.ai/docs/api-reference/overview#completionsresponse-format]
 */
@Serializable
public class OpenRouterChatCompletionStreamResponse(
    public val choices: List<OpenRouterStreamChoice>,
    override val created: Long,
    override val id: String,
    override val model: String,
    public val systemFingerprint: String? = null,
    @SerialName("object")
    public val objectType: String = "chat.completion.chunk",
    public val usage: OpenAIUsage? = null,
) : OpenAIBaseLLMStreamResponse

/**
 * OpenRouter API Error Response
 */
@Serializable
public class OpenRouterError(
    public val message: String,
    public val type: String? = null,
    public val code: String? = null,
)

internal object OpenRouterChatCompletionRequestSerializer :
    AdditionalPropertiesFlatteningSerializer<OpenRouterChatCompletionRequest>(OpenRouterChatCompletionRequest.serializer())
